import click
from pathlib import Path
from src.pipeline import Pipeline, configs, preprocess_configs

@click.group()
def cli():
    """Pipeline command line interface for processing PDF reports and questions."""
    pass

@cli.command()
@click.option('--parser', type=click.Choice(['docling', 'mineru']), default='docling', help='Parser to download models for')
def download_models(parser):
    """Download required models for the specified parser."""
    if parser == 'docling':
        click.echo("Downloading Docling models...")
        Pipeline.download_docling_models()
    elif parser == 'mineru':
        click.echo("MinerU doesn't require model downloads - it uses cloud API.")
    else:
        click.echo(f"Unknown parser: {parser}")

@cli.command()
@click.option('--parser', type=click.Choice(['docling', 'mineru']), default='mineru', help='PDF parser to use')
@click.option('--parallel/--sequential', default=False, help='Run parsing in parallel or sequential mode (Docling only)')
@click.option('--chunk-size', default=1, help='Number of PDFs to process in each worker (Docling only)')
@click.option('--max-workers', default=1, help='Number of parallel worker processes (Docling only)')
@click.option('--files', multiple=True, help='Specific PDF file names to parse (MinerU only). Provide multiple --files arguments.')
def parse_pdfs(parser, parallel, chunk_size, max_workers, files):
    """Parse PDF reports using Docling or MinerU.

    - For MinerU: if --files provided, only those files (by name) under pdf_reports will be parsed to markdown.
    - For Docling: current implementation ignores --files and processes all PDFs.
    """
    root_path = Path.cwd()
    
    # 根据选择的解析器创建配置
    if parser == 'docling':
        from src.pipeline import max_config_docling
        run_config = max_config_docling
        run_config.docling_parallel = parallel
        run_config.docling_chunk_size = chunk_size
        run_config.docling_max_workers = max_workers
    else:
        from src.pipeline import max_config_mineru
        run_config = max_config_mineru
    
    pipeline = Pipeline(root_path, run_config=run_config)
    
    click.echo(f"Parsing PDFs using {parser.upper()} parser...")
    if parser == 'docling':
        click.echo(f"Docling settings: parallel={parallel}, chunk_size={chunk_size}, max_workers={max_workers}")
    
    if parser == 'docling':
        # Docling 解析（并行/串行均复用并行接口，通过参数控制）
        pipeline.parse_pdf_reports_parallel(chunk_size=chunk_size, max_workers=max_workers)
    else:
        # MinerU 解析：若指定 --files，仅解析这些文件；否则解析全部
        if files:
            click.echo(f"MinerU will parse specific files: {', '.join(files)}")
            # 直接调用 Pipeline 内部的导出方法以逐个解析
            selected = list(files)
            success = 0
            for fname in selected:
                ok = pipeline.export_reports_to_markdown_mineru(fname)
                success += 1 if ok else 0
            click.echo(f"MinerU parsing finished. Success: {success}/{len(selected)}")
        else:
            pipeline.parse_pdf_reports_mineru()

@cli.command()
@click.option('--max-workers', default=10, help='Number of workers for table serialization')
def serialize_tables(max_workers):
    """Serialize tables in parsed reports using parallel threading (Docling only)."""
    root_path = Path.cwd()
    pipeline = Pipeline(root_path)
    
    click.echo(f"Serializing tables (max_workers={max_workers})...")
    pipeline.serialize_tables(max_workers=max_workers)

@cli.command()
def create_bm25_db():
    """Create BM25 database from chunked reports."""
    root_path = Path.cwd()
    pipeline = Pipeline(root_path)
    
    click.echo("Creating BM25 database...")
    pipeline.create_bm25_db()
    click.echo("BM25 database created successfully!")

@cli.command()
@click.option('--parser', type=click.Choice(['docling', 'mineru']), default='mineru', help='PDF parser to use')
@click.option('--parallel/--sequential', default=True, help='Run parsing in parallel or sequential mode (Docling only)')
@click.option('--chunk-size', default=2, help='Number of PDFs to process in each worker (Docling only)')
@click.option('--max-workers', default=10, help='Number of parallel worker processes (Docling only)')
def process_all_pdfs(parser, parallel, chunk_size, max_workers):
    """Complete PDF processing pipeline: parse -> chunk -> create vector DB."""
    root_path = Path.cwd()
    
    # 根据选择的解析器创建配置
    if parser == 'docling':
        from src.pipeline import max_config_docling
        run_config = max_config_docling
        run_config.docling_parallel = parallel
        run_config.docling_chunk_size = chunk_size
        run_config.docling_max_workers = max_workers
    else:
        from src.pipeline import max_config_mineru
        run_config = max_config_mineru
    
    pipeline = Pipeline(root_path, run_config=run_config)
    
    click.echo(f"Running complete PDF processing pipeline using {parser.upper()} parser...")
    if parser == 'docling':
        click.echo(f"Docling settings: parallel={parallel}, chunk_size={chunk_size}, max_workers={max_workers}")
    
    pipeline.process_all_pdfs(parser_type=parser)

@cli.command()
@click.option('--config', type=click.Choice(['ser_tab', 'no_ser_tab']), default='no_ser_tab', help='Configuration preset to use')
@click.option('--parser', type=click.Choice(['docling', 'mineru']), default='mineru', help='PDF parser that was used')
def process_reports(config, parser):
    """Process parsed reports through the pipeline stages."""
    root_path = Path.cwd()
    run_config = preprocess_configs[config]
    run_config.pdf_parser = parser  # 设置解析器类型
    pipeline = Pipeline(root_path, run_config=run_config)
    
    click.echo(f"Processing parsed reports (config={config}, parser={parser})...")
    pipeline.process_parsed_reports()

@cli.command()
@click.option('--config', type=click.Choice(['base', 'pdr', 'max', 'max_mineru', 'max_docling', 'max_mineru_enhanced', 'max_docling_enhanced']), default='max', help='Configuration preset to use')
@click.option('--retrieval-method', type=click.Choice(['vector', 'hybrid', 'enhanced_hybrid']), help='Override retrieval method (vector/hybrid/enhanced_hybrid)')
def process_questions(config, retrieval_method):
    """Process questions using the pipeline."""
    root_path = Path.cwd()
    run_config = configs[config]
    
    # 如果指定了检索方式，覆盖配置
    if retrieval_method:
        run_config.retrieval_method = retrieval_method
        if retrieval_method == "enhanced_hybrid":
            run_config.use_bm25_db = True
        click.echo(f"Override retrieval method to: {retrieval_method}")
    
    pipeline = Pipeline(root_path, run_config=run_config)
    
    click.echo(f"Processing questions (config={config}, parser={run_config.pdf_parser}, retrieval={run_config.retrieval_method})...")
    pipeline.process_questions()

if __name__ == '__main__':
    cli()